import pandas as pd
from sklearn.preprocessing import LabelEncoder, OneHotEncoder, StandardScaler, PolynomialFeatures
from sklearn.feature_selection import SelectKBest, f_regression
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

# 读取本地文件
data = pd.read_csv("C:/Users/wwj/Desktop/transactions.csv")

# 处理缺失值
data = data.fillna(data.mean())

# 处理分类特征
le = LabelEncoder()
data["category"] = le.fit_transform(data["category"])

# 处理数值特征
scaler = StandardScaler()
data[["age", "income"]] = scaler.fit_transform(data[["age", "income"]])

# 处理分类特征（独热编码）
ohe = OneHotEncoder()
category_ohe = ohe.fit_transform(data[["category"]])
category_ohe_df = pd.DataFrame(category_ohe.toarray(), columns=ohe.get_feature_names(["category"]))
data = pd.concat([data, category_ohe_df], axis=1)
data = data.drop("category", axis=1)

# 处理数值特征（多项式特征）
poly = PolynomialFeatures(degree=2)
age_income_poly = poly.fit_transform(data[["age", "income"]])
age_income_poly_df = pd.DataFrame(age_income_poly, columns=poly.get_feature_names(["age", "income"]))
data = pd.concat([data, age_income_poly_df], axis=1)
data = data.drop(["age", "income"], axis=1)

# 特征选择
selector = SelectKBest(f_regression, k=5)
X = data.drop("target", axis=1)
y = data["target"]
X_new = selector.fit_transform(X, y)
selected_features = X.columns[selector.get_support()]
data = data[selected_features]

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_new, y, test_size=0.2, random_state=42)

# 训练模型
model = LinearRegression()
model.fit(X_train, y_train)

# 预测测试集
y_pred = model.predict(X_test)

# 评估模型
mse = mean_squared_error(y_test, y_pred)
print("均方误差：", mse)